/*
 * @Author: Sian Li lisian_magic@163.com
 * @Date: 2024-08-21 17:10:05
 * @LastEditors: Sian Li lisian_magic@163.com
 * @LastEditTime: 2024-08-21 17:26:50
 * @FilePath: /Yolov5Detect/Detectors/src/Detector.cpp
 * @Description: 这是默认设置,请设置`customMade`, 打开koroFileHeader查看配置 进行设置: https://github.com/OBKoro1/koro1FileHeader/wiki/%E9%85%8D%E7%BD%AE
 */
#include "./Detector.h"
void makeRectSafe(Rect &rect, Mat &src)
{
    rect &= Rect(0, 0, src.cols, src.rows);
}

ArmorDetector::ArmorDetector(string engine_path, string onnx_path)
{
    assert(!engine_path.empty() && !onnx_path.empty());
    this->TensorRTEnginePath = engine_path;
    this->OnnxPath = onnx_path;
}

ArmorDetector::~ArmorDetector()
{
}

void ArmorDetector::accessModelTest()
{
    if (access(TensorRTEnginePath.c_str(), F_OK) != 0)
    {
#ifndef UseOneLayerInfer
        auto engine = this->armorTensorRT.createEngine(OnnxPath, 64, 640, 640, 0);
#else
        auto engine = this->armorTensorRT.createEngine(OnnxPath, 64, 1280, 1280, 0);
#endif
        this->armorTensorRT.saveEngineFile(engine, TensorRTEnginePath);
        delete engine;
    }
}

bool ArmorDetector::initModel()
{

    bool check = this->armorTensorRT.initModule(TensorRTEnginePath, 16, 12);

    return check;
}

vector<bboxAndRect> ArmorDetector::infer(Mat &image, vector<DetectBox> &targets)
{
    vector<vector<TRTInferV1::DetectionObj>> results_pre;
    if (targets.size() == 0)
    {
#ifdef ExperimentalOutput
        this->detectedArmorNumThisFrame = 0;
        this->averageConfThisFrame = 0.0;
#endif
        return {};
    }
    vector<Mat> preProcessedImage = this->preProcess(image, targets);
    results_pre = this->armorTensorRT.doInference(preProcessedImage, 0.1, 0.25, 0.45);
    this->reBuildBoxs(results_pre, targets, preProcessedImage);
#ifdef ExperimentalOutput
    this->detectedArmorNumThisFrame = this->results.size();
    this->averageConfThisFrame = sumConfAverage(this->results);
#endif
    return this->results;
}

vector<Mat> ArmorDetector::preProcess(Mat &image, vector<DetectBox> &movingTargets)
{
    vector<Mat> output;
    for (vector<DetectBox>::iterator it = movingTargets.begin(); it != movingTargets.end();)
    {
        if (it->x2 - it->x1 == 0 || it->y2 - it->y1 == 0)
        {
            it = movingTargets.erase(it);
        }
        else
        {
            Rect target = Rect(it->x1, it->y1, it->x2 - it->x1, it->y2 - it->y1);
            makeRectSafe(target, image);
            output.emplace_back(image(target).clone());
            ++it;
        }
    }

    return output;
}

void ArmorDetector::reBuildBoxs(vector<vector<TRTInferV1::DetectionObj>> &armors, vector<DetectBox> &boxs, vector<Mat> &img)
{
    vector<bboxAndRect>().swap(this->results);
    if (armors.size() != boxs.size())
        return;
    for (size_t i = 0; i < boxs.size(); ++i)
    {
        for (auto &it : armors[i])
        {
            this->results.emplace_back(bboxAndRect{ArmorBoundingBox{true,
                                                                    (float)it.x1 + boxs[i].x1,
                                                                    (float)it.y1 + boxs[i].y1,
                                                                    abs((float)(it.x2 - it.x1)),
                                                                    abs((float)(it.y2 - it.y1)),
                                                                    (float)it.classId, it.confidence},
                                                   boxs[i]});
            // this->logger->info("Arrmor: [x0] " + to_string(this->results.back().armor.x0) +
            //                    " [y0] " + to_string(this->results.back().armor.y0) +
            //                    " [w] " + to_string(this->results.back().armor.w) +
            //                    " [h] " + to_string(this->results.back().armor.h) +
            //                    " [cls] " + to_string(it.classId) +
            //                    " [conf] " + to_string(it.confidence));
        }
    }
}

void ArmorDetector::unInit()
{
    this->armorTensorRT.unInitModule();
}

CarDetector::CarDetector(string engine_path, string onnx_path)
{
    assert(!engine_path.empty() && !onnx_path.empty());
    this->TensorRTEnginePath = engine_path;
    this->OnnxPath = onnx_path;
}

CarDetector::~CarDetector()
{
}

void CarDetector::accessModelTest()
{
    if (access(TensorRTEnginePath.c_str(), F_OK) != 0)
    {
        auto engine = this->carTensorRT.createEngine(OnnxPath, 4, 1280, 1280, 0);
        this->carTensorRT.saveEngineFile(engine, TensorRTEnginePath);
        delete engine;
    }
}

bool CarDetector::initModel()
{

    cout << "CarDetector init Moudel" << endl;
    bool check = this->carTensorRT.initModule(TensorRTEnginePath, 1, 1);
    if (check)
        cout << "CarDetector Moudel inited" << endl;
    return check;
}

vector<DetectBox> CarDetector::infer(Mat &image)
{

    vector<vector<TRTInferV1::DetectionObj>> results;
    vector<Mat> srcs;
    srcs.emplace_back(image);
    results = this->carTensorRT.doInference(srcs, 0.1, 0.45, 0.3);
    vector<DetectBox> final_results;
    if (results.size() == 0)
        return final_results;
    for (size_t j = 0; j < results[0].size(); j++)
    {

        // cout << ("Car: [x1] " + to_string(results[0][j].x1) +
        //          " [y1] " + to_string(results[0][j].y1) +
        //          " [x2] " + to_string(results[0][j].x2) +
        //          " [y2] " + to_string(results[0][j].y2) +
        //          " [cls] " + to_string(results[0][j].classId) +
        //          " [conf] " + to_string(results[0][j].confidence))
        //      << endl;
        final_results.emplace_back(results[0][j].x1, results[0][j].y1, results[0][j].x2, results[0][j].y2, results[0][j].confidence, results[0][j].classId);
    }
    return final_results;
}

void CarDetector::unInit()
{
    this->carTensorRT.unInitModule();
}